import fitz  # PyMuPDF
import numpy as np
import pytesseract
from pdf2image import convert_from_path  # 如果需要，可以替换成其他方式的OCR
from PIL import Image
from concurrent.futures import ThreadPoolExecutor, as_completed

def extract_text_from_orange_pages(pdf_path, output_file=None):
    """
    提取PDF中橙色背景页面上的文字，支持扫描图像页面的OCR
    
    参数:
        pdf_path: PDF文件路径
        output_file: 输出文本文件路径(可选)
    
    返回:
        Dict[int, List[str]]: 按页码组织的提取文本字典
    """
    doc = fitz.open(pdf_path)
    results = {}
    
    # 目标背景的RGB范围 (195, 98, 4)
    ORANGE_RANGE = {
        "r": (195, 195),
        "g": (98, 98),
        "b": (4, 4)
    }

    # 提取时使用更严格的 OCR 引擎模式
    custom_config = r'--oem 3 --psm 6'  # OEM 3: 使用标准OCR引擎；PSM 6: 假设图像中有一块单一的文本区域

    
    # 背景检测区域 (整个页面的80%)
    BG_AREA_PERCENT = 0.8

    def process_page(page_num):
        page = doc.load_page(page_num)
        
        # 创建页面的高分辨率图像(150 DPI 提升速度)
        pix = page.get_pixmap(matrix=fitz.Matrix(150/72, 150/72), colorspace="rgb")
        
        # 将图像转换为numpy数组
        img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, 3)
        
        # 分析背景颜色 - 取页面中心区域
        height, width = img_array.shape[:2]
        x1, y1 = int(width * (1 - BG_AREA_PERCENT)/2), int(height * (1 - BG_AREA_PERCENT)/2)
        x2, y2 = int(width * (1 + BG_AREA_PERCENT)/2), int(height * (1 + BG_AREA_PERCENT)/2)
        center_area = img_array[y1:y2, x1:x2]
        
        # 计算中心区域的主要颜色
        flat_area = center_area.reshape(-1, 3)
        r, g, b = np.median(flat_area, axis=0)
        
        # 检查是否是目标背景颜色
        is_orange = (ORANGE_RANGE["r"][0] <= r <= ORANGE_RANGE["r"][1] and
                     ORANGE_RANGE["g"][0] <= g <= ORANGE_RANGE["g"][1] and
                     ORANGE_RANGE["b"][0] <= b <= ORANGE_RANGE["b"][1])
        
        if is_orange:
            print(f"检测到橙色背景 - 页面 {page_num+1} - RGB: ({r}, {g}, {b})")
            # 使用OCR提取图像中的文本
            page_image = Image.fromarray(img_array)
            extracted_text = pytesseract.image_to_string(page_image)
            
            # 如果提取的文本非空
            if extracted_text.strip():
                return (page_num + 1, [extracted_text.strip()])
        return None
    
    # 使用线程池并行处理页面
    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = [executor.submit(process_page, page_num) for page_num in range(doc.page_count)]
        for future in as_completed(futures):
            result = future.result()
            if result:
                page_num, texts = result
                results[page_num] = texts
                print(f"提取到 {len(texts)} 条文字，页面 {page_num} 内容: {texts[0]}")
    
    # 保存结果到文件
    if output_file and results:
        with open(output_file, "w", encoding="utf-8") as f:
            for pg, texts in results.items():
                f.write(f"{texts[0]} — {pg}\n")
        print(f"\n结果已保存到: {output_file}")
    
    doc.close()
    return results

# 使用示例
if __name__ == "__main__":
    # 替换为你的PDF文件路径
    pdf_file = "C:\\Users\\Rusz\\Downloads\\bkqs_all.pdf"
    
    # 执行提取
    extracted = extract_text_from_orange_pages(
        pdf_file,
        output_file="C:\\Users\\Rusz\\Downloads\\pdf_index.txt"
    )
    
    # 打印结果摘要
    if extracted:
        print("\n提取摘要:")
        for pg, texts in extracted.items():
            print(f" - 页面 {pg}: {len(texts)} 条文字")
    else:
        print("未找到含有目标背景颜色的页面")
